import math
import random


def rank(sequence, value):
    count = 0
    for each in sequence:
        if each < value:
            count += 1
    return count + 1


def subset_by_range(s, a, b):
    subset = []
    for each in s:
        if a <= each <= b:
            subset.append(each)
    return subset


def _lazyselect(s, k, r_exp=3 / 4):
    n = len(s)
    r_size = n ** r_exp
    r = random.choices(s, k=math.floor(r_size))
    r.sort()
    x = k * n ** (r_exp - 1)
    l = max(math.floor(x - math.sqrt(n)), 1)
    h = min(math.floor(x + math.sqrt(n)), len(r))
    a, b = r[l - 1], r[h - 1]
    rank_s_a, rank_s_b = rank(s, a), rank(s, b)
    if k < rank_s_a or k > rank_s_b:
        # print('k < rank_s_a or k > rank_s_b')
        # return None
        return 'k-th element is not in P'
    p = subset_by_range(s, a, b)
    if len(p) >= 4 * r_size:
        # print('len(p) >= 4 * r_size')
        # return None
        # return 'len(p) >= 4 * r_size'
        return '|p| >= 4 * n ^ sampling_ratio'
    p.sort()
    return p[k - rank_s_a]


def lazyselect_once(sequence, k, r_exp=3 / 4):
    if 0 < k <= len(sequence):
        return _lazyselect(sequence, k, r_exp=r_exp)
    else:
        raise ValueError("Invalid k")


def lazyselect(sequence, k):
    if not 0 < k <= len(sequence):
        raise ValueError("Invalid k")

    ret = _lazyselect(sequence, k)
    while ret is None:
        ret = _lazyselect(sequence, k)
    return ret


if __name__ == '__main__':
    s = list(range(100))
    random.shuffle(s)
    print(s)
    _k = 1
    print(f'k = {_k}:', _lazyselect(s, _k))
